import torch.nn as nn
import torch.nn.functional as F


class FPN(nn.Module):
    def __init__(self):
        super(FPN, self).__init__()

        # 1✖1卷积：修改至与上采样结构具有相同的channels，为特征融合做准备
        self.lateral_conv1 = nn.Conv2d(64, 128, kernel_size=1)
        self.lateral_conv2 = nn.Conv2d(128, 256, kernel_size=1)
        self.lateral_conv3 = nn.Conv2d(256, 512, kernel_size=1)
        self.lateral_conv4 = nn.Conv2d(128, 64, kernel_size=1)
        # 3✖3卷积：融合之后重新卷积，消除混沌
        self.smooth_conv2 = nn.Conv2d(256, 128, kernel_size=3, stride=1, padding=1)
        self.smooth_conv3 = nn.Conv2d(512, 256, kernel_size=3, stride=1, padding=1)
        self.smooth_conv4 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)

    def up_sample(self, x, y):
        # 提取采样尺寸
        b, c, h, w = y.size()
        # 上采样
        x = F.interpolate(x, size=(h, w), mode='bilinear', align_corners=True)
        # 融合
        output = x + y
        return output

    def forward(self, x2, x3, x4):
        c2 = self.lateral_conv2(x2)
        c3 = self.lateral_conv3(x3)
        s4 = self.up_sample(x4, c3)
        p3 = c3 + s4
        s3 = self.up_sample(x3, c2)
        p2 = c2 + s3
        out_p2 = self.smooth_conv2(p2)
        out_p3 = self.smooth_conv3(p3)
        out_p4 = self.smooth_conv4(x4)

        return out_p2, out_p3, out_p4
